#!/usr/bin/env python
# encoding=utf-8

import json
import requests
import sys

import tornado.ioloop
import tornado.web
from cassandra import ConsistencyLevel
from cassandra.cluster import Cluster
from cassandra.policies import RoundRobinPolicy
from concurrent.futures import ThreadPoolExecutor
from pymongo import MongoClient
from scpy.logger import get_logger
from tornado import gen
from tornado.concurrent import Future, run_on_executor
from tornado.ioloop import IOLoop
from tornado.web import HTTPError
from xtls.timeparser import now
from xtls.codehelper import trytry

from config import *

reload(sys)
sys.setdefaultencoding('utf-8')
logger = get_logger(__file__)
MONGO = MongoClient(MONGO_HOST, MONGO_PORT)

CLUSTER = Cluster(CASSANDRA_CONTACT_POINTS, load_balancing_policy=RoundRobinPolicy(), protocol_version=2)
SESSION = CLUSTER.connect(KEYSPACE)
EXPIRE = 86400 * 7
SQL = u'INSERT INTO news_by_company_name (company_name,period,url,content,time) VALUES (?,?,?,?,?)'
INSERT_STMT = SESSION.prepare(SQL)

QUERYS = {
    'news': SESSION.prepare("SELECT content FROM news_by_company_name WHERE company_name = ? and period IN ?"),
}

for query in QUERYS.values():
    query.consistency_level = ConsistencyLevel.LOCAL_QUORUM


def save(company, news):
    period = 2016
    stmt = INSERT_STMT.bind([company, period, news['url'], json.dumps(news, ensure_ascii=False), now()])
    SESSION.execute(stmt)


def wrap_cassandra_execute(io_loop, cql, parmeters=None, trace=False):
    tornado_future = Future()
    cassandra_future = SESSION.execute_async(cql, parmeters, trace)

    def _result(result):
        io_loop.add_callback(tornado_future.set_result, result)

    def _error(err):
        io_loop.add_callback(tornado_future.set_exception, err)

    cassandra_future.add_callbacks(_result, _error)
    return tornado_future


class NewsHandler(tornado.web.RequestHandler):
    executor = ThreadPoolExecutor(max_workers=4)

    @gen.coroutine
    def get(self):
        self.set_header('Content-Type', 'application/json; charset="utf-8"')
        company_name = self.get_argument('companyName')
        call_back = self.get_argument('callback', '')
        update_anyway = self.get_argument('update', 'N')

        done = MONGO[DB][COLL_STATUS].find_one({'_id': company_name})
        if not done:
            MONGO[DB][COLL_STATUS].insert_one({
                '_id': company_name,
                'callback': call_back,
                'updateTime': now(),
                'done': []
            })
            raise HTTPError(204, '')

        if not done['done']:
            raise HTTPError(204, '')

        if (now() - done['updateTime']).seconds > EXPIRE or update_anyway == 'Y':
            MONGO[DB][COLL_STATUS].update_one({'_id': company_name}, {
                '_id': company_name,
                'callback': call_back,
                'updateTime': now(),
                'done': []
            })
        stmt = QUERYS['news'].bind((company_name, [2016, 201511, 201512, 201601, 201602]))
        row = yield wrap_cassandra_execute(IOLoop.current(), stmt)
        if row:
            news_list = [json.loads(item.content) for item in row]
            if news_list:
                self.write(json.dumps(news_list))
                return
        raise HTTPError(404, 'not found')


class PullJobHandler(tornado.web.RequestHandler):
    executor = ThreadPoolExecutor(max_workers=4)

    @run_on_executor
    def get(self):
        tp = self.get_argument('type', 'wechat')
        item = MONGO[DB][COLL_STATUS].find_one({'done': {'$nin': [tp]}})
        if item:
            self.write(item['_id'])
        else:
            raise HTTPError(404, 'no more jobs.')


class ReportHandler(tornado.web.RequestHandler):
    executor = ThreadPoolExecutor(max_workers=4)

    @run_on_executor
    def post(self):
        company_name = self.get_argument('companyName')
        tp = self.get_argument('type', 'wechat')
        data = self.get_argument('data')

        item = MONGO[DB][COLL_STATUS].find_one({'_id': company_name})
        done = item['done'] + [tp]
        MONGO[DB][COLL_STATUS].update_one({'_id': company_name}, {'$set': {'done': done}})

        if data:
            for news in json.loads(data, encoding='utf-8'):
                save(company_name, news)

        with trytry():
            if item['callback']:
                requests.post(item['callback'], data={'companyName': company_name, 'data': data})

        self.write('{"msg":"done"}')


if __name__ == '__main__':
    SERVER_PORT = 56789 if len(sys.argv) < 2 else int(sys.argv[1])
    application = tornado.web.Application([
        (r'/news', NewsHandler),
        (r'/report', ReportHandler),
        (r'/pullJob', PullJobHandler),
    ], debug=False)
    application.listen(SERVER_PORT, address='0.0.0.0')
    logger.info('tornado server started on port %s', SERVER_PORT)
    tornado.ioloop.IOLoop.instance().start()
